"""
测试数据预处理模块
"""
import os
import sys
import unittest
import numpy as np

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))

from sichuanmajiang.data.preprocessor import DataPreprocessor, FeatureEngineer
from sichuanmajiang.tests.test_utils import create_test_game_state, create_test_hand, create_test_directory, remove_test_directory


class TestDataPreprocessor(unittest.TestCase):
    """
    测试DataPreprocessor类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_data')
        create_test_directory(self.test_dir)
        
        # 创建预处理器
        self.preprocessor = DataPreprocessor()
        
        # 创建测试数据
        self.raw_data = []
        for i in range(10):
            game_state = create_test_game_state()
            game_state['game_round'] = i
            game_state['timestamp'] = i * 1000
            self.raw_data.append({
                'state': game_state,
                'action': i % 6,
                'reward': float(i)
            })
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_clean_data(self):
        """
        测试数据清洗
        """
        # 添加一些异常数据
        test_data = self.raw_data.copy()
        test_data.append({
            'state': None,  # 无效的状态
            'action': 0,
            'reward': 0.0
        })
        test_data.append({
            'state': create_test_game_state(),
            'action': -1,  # 无效的动作
            'reward': 0.0
        })
        
        # 清洗数据
        cleaned_data = self.preprocessor.clean_data(test_data)
        
        # 检查清洗后的数据长度（应该过滤掉2个异常数据）
        self.assertEqual(len(cleaned_data), len(self.raw_data))
    
    def test_extract_features(self):
        """
        测试特征提取
        """
        # 提取特征
        features = self.preprocessor.extract_features(self.raw_data)
        
        # 检查特征是否被提取
        self.assertIn('states', features)
        self.assertIn('actions', features)
        self.assertIn('rewards', features)
        
        # 检查长度
        self.assertEqual(len(features['states']), len(self.raw_data))
        self.assertEqual(len(features['actions']), len(self.raw_data))
        self.assertEqual(len(features['rewards']), len(self.raw_data))
    
    def test_normalize_features(self):
        """
        测试特征归一化
        """
        # 创建一些测试特征
        test_features = np.array([
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0],
            [7.0, 8.0, 9.0]
        ], dtype=np.float32)
        
        # 归一化特征
        normalized_features = self.preprocessor.normalize_features(test_features)
        
        # 检查是否归一化（均值应接近0，标准差应接近1）
        self.assertAlmostEqual(normalized_features.mean(), 0.0, places=6)
        self.assertAlmostEqual(normalized_features.std(), 1.0, places=6)
    
    def test_split_data(self):
        """
        测试数据划分
        """
        # 创建一些测试数据
        test_states = np.random.rand(100, 10)
        test_actions = np.random.randint(0, 6, 100)
        test_rewards = np.random.rand(100)
        
        # 划分数据
        splits = self.preprocessor.split_data(
            test_states, test_actions, test_rewards,
            train_ratio=0.7, val_ratio=0.15, test_ratio=0.15
        )
        
        # 检查划分比例
        self.assertEqual(len(splits['train_states']), 70)
        self.assertEqual(len(splits['val_states']), 15)
        self.assertEqual(len(splits['test_states']), 15)
    
    def test_save_preprocessor(self):
        """
        测试保存预处理器
        """
        # 先训练预处理器
        test_features = np.random.rand(100, 10)
        self.preprocessor.normalize_features(test_features)  # 这会更新预处理器的统计信息
        
        # 保存预处理器
        save_path = os.path.join(self.test_dir, 'preprocessor.pkl')
        self.preprocessor.save_preprocessor(save_path)
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(save_path))
    
    def test_load_preprocessor(self):
        """
        测试加载预处理器
        """
        # 先训练并保存预处理器
        test_features = np.random.rand(100, 10)
        self.preprocessor.normalize_features(test_features)
        
        save_path = os.path.join(self.test_dir, 'preprocessor.pkl')
        self.preprocessor.save_preprocessor(save_path)
        
        # 创建新的预处理器并加载
        new_preprocessor = DataPreprocessor()
        new_preprocessor.load_preprocessor(save_path)
        
        # 检查是否加载成功（比较均值）
        self.assertTrue(hasattr(new_preprocessor, 'feature_mean'))


class TestFeatureEngineer(unittest.TestCase):
    """
    测试FeatureEngineer类
    """
    def setUp(self):
        """
        测试前的设置
        """
        # 创建特征工程师
        self.engineer = FeatureEngineer()
        
        # 创建测试手牌
        self.test_hand = create_test_hand()
    
    def test_encode_hand(self):
        """
        测试手牌编码
        """
        # 编码手牌
        encoded_hand = self.engineer.encode_hand(self.test_hand)
        
        # 检查编码结果
        self.assertEqual(len(encoded_hand), 108)  # 四川麻将最多108种牌
        
        # 检查是否正确编码了刻子
        self.assertEqual(encoded_hand[0], 3)  # 假设1对应索引0
        self.assertEqual(encoded_hand[1], 3)  # 假设2对应索引1
        self.assertEqual(encoded_hand[2], 3)  # 假设3对应索引2
    
    def test_create_board_representation(self):
        """
        测试创建棋盘表示
        """
        # 创建游戏状态
        game_state = create_test_game_state()
        
        # 创建棋盘表示
        board_repr = self.engineer.create_board_representation(game_state)
        
        # 检查棋盘表示的维度
        self.assertEqual(len(board_repr.shape), 3)  # 应该是3D张量
    
    def test_calculate_tile_frequency(self):
        """
        测试计算牌的频率
        """
        # 计算频率
        frequency = self.engineer.calculate_tile_frequency(self.test_hand)
        
        # 检查频率计算
        self.assertEqual(len(frequency), 108)
        self.assertEqual(frequency[0], 3)  # 1出现了3次
    
    def test_calculate_meld_features(self):
        """
        测试计算组合特征
        """
        # 计算组合特征
        meld_features = self.engineer.calculate_meld_features(self.test_hand)
        
        # 检查组合特征
        self.assertIn('triplets', meld_features)
        self.assertIn('sequences', meld_features)
        self.assertIn('pairs', meld_features)
        
        # 检查刻子数量（应该有3个刻子）
        self.assertEqual(meld_features['triplets'], 3)
    
    def test_extract_hand_features(self):
        """
        测试提取手牌特征
        """
        # 提取手牌特征
        hand_features = self.engineer.extract_hand_features(self.test_hand)
        
        # 检查特征长度
        self.assertGreater(len(hand_features), 0)
    
    def test_create_feature_vector(self):
        """
        测试创建特征向量
        """
        # 创建游戏状态
        game_state = create_test_game_state()
        
        # 创建特征向量
        feature_vector = self.engineer.create_feature_vector(game_state)
        
        # 检查特征向量
        self.assertIsInstance(feature_vector, np.ndarray)
        self.assertEqual(len(feature_vector.shape), 1)  # 应该是1D数组


if __name__ == '__main__':
    unittest.main()
